#include "StdAfx.h"

void
CPrgExeState::PrintHelp()
{
	wprintf(L"Usage: %s Capture EventNumber [m_LogFile] [-options]\n", AppName);
	wprintf(L"  Capture     - Name of capture file to use.  use -o to overwrite if capture already exists.\n");
	wprintf(L"  EventNumber - numeric event error message to stop on.\n");
	wprintf(L"  LogFile     - For example, Application, Security, System.  Default searches all logs.\n");
	wprintf(L"\n");
	wprintf(L"Options:\n");
	wprintf(L"  -b #     - Buffer size in Mbytes for capture.  Default is 100MB.\n");
	wprintf(L"  -c       - Use chain capture instead of the default of circular.\n");
	wprintf(L"  -f       - Filter to use for capturing traffic.\n");
	wprintf(L"  -o       - Overwrite capture if it exists.\n");
	wprintf(L"  -d       - Disable Conversations.  Warning, you could shoot yourself in the foot.\n");
	wprintf(L"  -n #     - Number of adapter to capture on.  Use Nmcap /displaynetworks to get list\n");
	wprintf(L"  -v       - Be verbose.  Show NPL compilation messages.\n");
}

//
// Initializes the class by collecting all the command line arguments and saving
// them in class variables.
//
CPrgExeState::CPrgExeState(int argc, WCHAR* argv[])
{
	AppName = argv[0];

    // We expect at least two args, a m_Filter, and then a file.
    if(argc <= 2){
        _IsValidCommandLine = false;
		return;
    }

	CapFileSize = 100000000; // 100 MB
	m_LogFile = NULL;
	ChainCapture = false;
	m_Filter = NULL;
	OverwriteCapture = false;
	ErrorLoadingNPL = false;
	AdapterNumber = -1;
	DisableConv = false;
	m_Verbose = false;

	bool GotFile = false;
	bool GotEventNum = false;
	for(int i = 1; i < argc; i++)
	{
		if(argv[i][0] == '-' || argv[i][0] == '/')
		{
			switch(argv[i][1])
			{
			case 'b':
				i++;
				if(i < argc && iswdigit(argv[i][0]))
				{
					CapFileSize = _wtol(argv[i]) * 1000000L;
				} 
				else
				{
					wprintf(L"-b option requires a number as an argument\n");
					_IsValidCommandLine = false;
					return;
				}
				break;
			case 'f':
				i++;
				if(i < argc)
				{
					m_Filter = argv[i];
				} 
				else
				{
					_IsValidCommandLine =  false;
					return;
				}
				break;
			case 'c':
				ChainCapture = true;
				break;
			case 'o':
				OverwriteCapture = true;
				break;
			case 'v':
				m_Verbose = true;
				break;
			case 'd':
				DisableConv = true;
				break;
			case 'n':
				i++;
				if(i < argc)
				{
					AdapterNumber = _wtol(argv[i]);
				}
				break;
			}
		} else {
			// First param is the capture file.
			if(!GotFile)
			{
				m_CapFile = argv[i];
				GotFile = true;
			} else 
			// Second is the Event Number
			if(!GotEventNum)
			{
				if(iswdigit(argv[i][0]))
				{
					m_EventNum = argv[i];
					GotEventNum = true;
				} else 
				{
					wprintf(L"m_EventNum should be a number.\n");
					_IsValidCommandLine =  false;
				}
			} else
			// Third is the Log file, but it doesn't have to be present.
			{
				m_LogFile = argv[i];
			}
		}
	}

	// If we've got both a file name and Event Number we are set.
	if(GotFile && GotEventNum)
	{
		_IsValidCommandLine =  true;
	}
	else
	{
		_IsValidCommandLine =  false;
	}
}

CPrgExeState::~CPrgExeState(void)
{
	NmCloseHandle(myCaptureEngine);
	NmCloseHandle(m_myCapFile);

	if(m_Filter != NULL)
	{
		NmCloseHandle(myFrameParser);
		NmCloseHandle(myFrameParserConfig);
		NmCloseHandle(myNplParser);
	}
}

// 
// Callback for building the NPL.  This is only needed if a m_Filter is used.
// Only errors will be printed out, unless m_Verbose mode is specified.
//
void __stdcall 
CPrgExeState::MyParserBuild(PVOID Context, ULONG StatusCode, LPCWSTR lpDescription, ULONG ErrorType)
{
	CPrgExeState *pesContext = (CPrgExeState *)Context;

	if(ErrorType == NmApiCallBackMsgTypeError)
	{
	    wprintf(L"%s\n", lpDescription);
		pesContext->ErrorLoadingNPL = true;
	} else {
		if(pesContext->m_Verbose)
		{
		    wprintf(L"%s\n", lpDescription);
		}
		pesContext->ErrorLoadingNPL = false;
	}
}

//
// Creates a frame parser and m_Filter if a m_Filter is specified.  Also disables
// conversations if that option is presented.  User can shot themselvs in the
// foot if a m_Filter requires converations and they disalble them.
//
void
CPrgExeState::LoadNPL(void)
{

    // Loads default NPL Path
    m_LastError = NmLoadNplParser(NULL, NmAppendRegisteredNplSets, (&CPrgExeState::MyParserBuild), this, &myNplParser);

	if(m_LastError != ERROR_SUCCESS)
    {
        wprintf(L"Unable to load NPL\n");
        ErrorLoadingNPL = true;
		return;
    }

    m_LastError = NmCreateFrameParserConfiguration(myNplParser, (&CPrgExeState::MyParserBuild), this, &myFrameParserConfig);
    if(m_LastError != ERROR_SUCCESS)
    {
        wprintf(L"Unable to load parser config\n");
        ErrorLoadingNPL = true;
		return;
    }

	m_LastError = NmAddFilter(myFrameParserConfig, m_Filter, &FilterID);
    if(m_LastError != ERROR_SUCCESS)
    {
        wprintf(L"Unable to add m_Filter %S\n", m_Filter);
        ErrorLoadingNPL = true;
		return;

    }

	if(!DisableConv)
	{
		NmConfigConversation(myFrameParserConfig, NmConversationOptionNone, TRUE);
	}
    
    myFrameParser = INVALID_HANDLE_VALUE;
    m_LastError = NmCreateFrameParser(myFrameParserConfig, &myFrameParser);
    if(ERROR_SUCCESS != m_LastError)
    {
        ErrorLoadingNPL = true;
    }
}

//
// Checks to see if the capture file is there if the option to overwrite capture
// is not specified.  Creates the capture file based on the options present for
// size and wrap vs. chained.
//
bool 
CPrgExeState::CreateCap(void)
{
	// See if there's already a caputre file there so we don't overwrite.

	if(!OverwriteCapture)
	{
		if(GetFileAttributes(m_CapFile) != INVALID_FILE_ATTRIBUTES)
		{
			wprintf(L"Capture file %s already exists\n", m_CapFile);
			return false;
		}
	}

	m_LastError = NmCreateCaptureFile(m_CapFile, CapFileSize, ChainCapture ? NmCaptureFileChain : NmCaptureFileWrapAround, &m_myCapFile, NULL);
    if(m_LastError != ERROR_SUCCESS)
    {
        wprintf(L"Error opening capture file, 0x%X\n", m_LastError);
        return false;
    }

	return true;
}

//
// Checks to see if a rawFrame passes a m_Filter.
//
BOOL
CPrgExeState::PassFilter(HANDLE rawFrame)
{
	if(m_Filter != NULL)
	{
		HANDLE ParsedFrame;

		NmParseFrame(myFrameParser, rawFrame, 0, 0, &ParsedFrame, NULL);

		BOOL passed = true;

		NmEvaluateFilter(ParsedFrame, FilterID, &passed);

		NmCloseHandle(ParsedFrame);

		return passed;
	}
	else
	{
		return true;
	}
}

//
// Callback for each frame indicated.  If a m_Filter is provided, it checks that the frame
// passes before storing in the frame.
//
void __stdcall 
CPrgExeState::MyFrameIndication(HANDLE hCapEng, ULONG ulAdaptIdx, PVOID pContext, HANDLE hRawFrame)
{
    CPrgExeState *pesContext = (CPrgExeState *)pContext;

	if(pesContext->PassFilter(hRawFrame))
	{
		NmAddFrame(pesContext->m_myCapFile, hRawFrame);
	}
}

//
// Initializes the capture engine and configures all adapters specified.
//
bool 
CPrgExeState::InitCapture(void)
{
    // Open capture engine.
    
    m_LastError = NmOpenCaptureEngine(&myCaptureEngine);
    if(m_LastError != ERROR_SUCCESS)
    {
        wprintf(L"Error opening capture engine, 0x%X\n", m_LastError);
        return false;
    }

	if(AdapterNumber < 0)
	{
		// Configure to capture on all adapters.

		NmGetAdapterCount(myCaptureEngine, &AdpCount);
		int NumAdpts = 0;

		for(DWORD n=0; n< AdpCount;n ++){
			m_LastError = NmConfigAdapter(myCaptureEngine, n, (&CPrgExeState::MyFrameIndication), this, NmReturnRemainFrames);
			if(m_LastError != ERROR_SUCCESS)
			{
				wprintf(L"Error configuration adapter %d, 0x%X\n", n, m_LastError);
			} else {
				NumAdpts++;
			}
		}

		if(NumAdpts <= 0)
		{
			wprintf(L"Unable to open any adapters\n");
			return false;
		}
	} else {
		m_LastError = NmConfigAdapter(myCaptureEngine, AdapterNumber, (&CPrgExeState::MyFrameIndication), this, NmReturnRemainFrames);
		if(m_LastError != ERROR_SUCCESS)
		{
			wprintf(L"Error configuration adapter %d, 0x%X\n", AdapterNumber, m_LastError);
			return false;
		}
	}

	return true;
}

// 
// Starts capture on all adapters specified.
//
bool 
CPrgExeState::StartCapture(void)
{
	if(AdapterNumber < 0)
	{
		for(DWORD n=0; n< AdpCount;n ++){
			m_LastError = NmStartCapture(myCaptureEngine, n, NmLocalOnly);
			if(m_LastError != ERROR_SUCCESS)
			{
				wprintf(L"Error 0x%X starting adapter #%d\n", m_LastError, n);
				return false;
			}
		}
	} else
	{
		m_LastError = NmStartCapture(myCaptureEngine, AdapterNumber, NmLocalOnly);
		if(m_LastError != ERROR_SUCCESS)
		{
			wprintf(L"Error 0x%X starting adapter #%d\n", m_LastError, AdapterNumber);
			return false;
		}
	}

	return true;
}

// 
// Stops capture on all adapters specified.
//
bool 
CPrgExeState::StopCapture(void)
{
	if(AdapterNumber < 0)
	{
		for(DWORD n=0; n< AdpCount;n ++){
			m_LastError = NmStopCapture(myCaptureEngine, n);
			if(m_LastError != ERROR_SUCCESS)
			{
				wprintf(L"Error 0x%X stopping adapter #%d\n", m_LastError, n);
				return false;
			}
		}
	} else
	{
		m_LastError = NmStopCapture(myCaptureEngine, AdapterNumber);
		if(m_LastError != ERROR_SUCCESS)
		{
			wprintf(L"Error 0x%X starting adapter #%d\n", m_LastError, AdapterNumber);
			return false;
		}
	}
	return true;
}
